﻿using System;
using System.Diagnostics;
using System.IO;
using System.Runtime.InteropServices;

namespace Artemis.UI.Windows.Utilities;

public static class ProcessUtilities
{
    public static Process? RunAsDesktopUser(string fileName, string arguments, bool hideWindow)
    {
        if (string.IsNullOrWhiteSpace(fileName))
            throw new ArgumentException("Value cannot be null or whitespace.", nameof(fileName));

        // To start process as shell user you will need to carry out these steps:
        // 1. Enable the SeIncreaseQuotaPrivilege in your current token
        // 2. Get an HWND representing the desktop shell (GetShellWindow)
        // 3. Get the Process ID(PID) of the process associated with that window(GetWindowThreadProcessId)
        // 4. Open that process(OpenProcess)
        // 5. Get the access token from that process (OpenProcessToken)
        // 6. Make a primary token with that token(DuplicateTokenEx)
        // 7. Start the new process with that primary token(CreateProcessWithTokenW)

        IntPtr hProcessToken = IntPtr.Zero;
        // Enable SeIncreaseQuotaPrivilege in this process.  (This won't work if current process is not elevated.)
        try
        {
            IntPtr process = GetCurrentProcess();
            if (!OpenProcessToken(process, 0x0020, ref hProcessToken))
                return null;

            TOKEN_PRIVILEGES tkp = new()
            {
                PrivilegeCount = 1,
                Privileges = new LUID_AND_ATTRIBUTES[1]
            };

            if (!LookupPrivilegeValue(null, "SeIncreaseQuotaPrivilege", ref tkp.Privileges[0].Luid))
                return null;

            tkp.Privileges[0].Attributes = 0x00000002;

            if (!AdjustTokenPrivileges(hProcessToken, false, ref tkp, 0, IntPtr.Zero, IntPtr.Zero))
                return null;
        }
        finally
        {
            CloseHandle(hProcessToken);
        }

        // Get an HWND representing the desktop shell.
        // CAVEATS:  This will fail if the shell is not running (crashed or terminated), or the default shell has been
        // replaced with a custom shell.  This also won't return what you probably want if Explorer has been terminated and
        // restarted elevated.
        IntPtr hwnd = GetShellWindow();
        if (hwnd == IntPtr.Zero)
            return null;

        IntPtr hShellProcess = IntPtr.Zero;
        IntPtr hShellProcessToken = IntPtr.Zero;
        IntPtr hPrimaryToken = IntPtr.Zero;
        try
        {
            // Get the PID of the desktop shell process.
            uint dwPID;
            if (GetWindowThreadProcessId(hwnd, out dwPID) == 0)
                return null;

            // Open the desktop shell process in order to query it (get the token)
            hShellProcess = OpenProcess(ProcessAccessFlags.QueryInformation, false, dwPID);
            if (hShellProcess == IntPtr.Zero)
                return null;

            // Get the process token of the desktop shell.
            if (!OpenProcessToken(hShellProcess, 0x0002, ref hShellProcessToken))
                return null;

            uint dwTokenRights = 395U;

            // Duplicate the shell's process token to get a primary token.
            // Based on experimentation, this is the minimal set of rights required for CreateProcessWithTokenW (contrary to current documentation).
            if (!DuplicateTokenEx(hShellProcessToken, dwTokenRights, IntPtr.Zero, SECURITY_IMPERSONATION_LEVEL.SecurityImpersonation, TOKEN_TYPE.TokenPrimary, out hPrimaryToken))
                return null;

            // Start the target process with the new token.
            STARTUPINFO si = new();
            if (hideWindow)
            {
                si.dwFlags = 0x00000001;
                si.wShowWindow = 0;
            }

            PROCESS_INFORMATION pi = new();
            if (!CreateProcessWithTokenW(hPrimaryToken, 0, fileName, $"\"{fileName}\" {arguments}", 0, IntPtr.Zero, Path.GetDirectoryName(fileName)!, ref si, out pi))
            {
                // Get the last error and display it.
                int error = Marshal.GetLastWin32Error();
                return null;
            }

            return Process.GetProcessById(pi.dwProcessId);
        }
        finally
        {
            CloseHandle(hShellProcessToken);
            CloseHandle(hPrimaryToken);
            CloseHandle(hShellProcess);
        }
    }

    #region Interop

    private struct TOKEN_PRIVILEGES
    {
        public uint PrivilegeCount;

        [MarshalAs(UnmanagedType.ByValArray, SizeConst = 1)]
        public LUID_AND_ATTRIBUTES[] Privileges;
    }

    [StructLayout(LayoutKind.Sequential, Pack = 4)]
    private struct LUID_AND_ATTRIBUTES
    {
        public LUID Luid;
        public uint Attributes;
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct LUID
    {
        public readonly uint LowPart;
        public readonly int HighPart;
    }

    [Flags]
    private enum ProcessAccessFlags : uint
    {
        All = 0x001F0FFF,
        Terminate = 0x00000001,
        CreateThread = 0x00000002,
        VirtualMemoryOperation = 0x00000008,
        VirtualMemoryRead = 0x00000010,
        VirtualMemoryWrite = 0x00000020,
        DuplicateHandle = 0x00000040,
        CreateProcess = 0x000000080,
        SetQuota = 0x00000100,
        SetInformation = 0x00000200,
        QueryInformation = 0x00000400,
        QueryLimitedInformation = 0x00001000,
        Synchronize = 0x00100000
    }

    private enum SECURITY_IMPERSONATION_LEVEL
    {
        SecurityAnonymous,
        SecurityIdentification,
        SecurityImpersonation,
        SecurityDelegation
    }

    private enum TOKEN_TYPE
    {
        TokenPrimary = 1,
        TokenImpersonation
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct PROCESS_INFORMATION
    {
        public readonly IntPtr hProcess;
        public readonly IntPtr hThread;
        public readonly int dwProcessId;
        public readonly int dwThreadId;
    }

    [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
    private struct STARTUPINFO
    {
        public readonly int cb;
        public readonly string lpReserved;
        public readonly string lpDesktop;
        public readonly string lpTitle;
        public readonly int dwX;
        public readonly int dwY;
        public readonly int dwXSize;
        public readonly int dwYSize;
        public readonly int dwXCountChars;
        public readonly int dwYCountChars;
        public readonly int dwFillAttribute;
        public int dwFlags;
        public short wShowWindow;
        public readonly short cbReserved2;
        public readonly IntPtr lpReserved2;
        public readonly IntPtr hStdInput;
        public readonly IntPtr hStdOutput;
        public readonly IntPtr hStdError;
    }

    [DllImport("kernel32.dll", ExactSpelling = true)]
    private static extern IntPtr GetCurrentProcess();

    [DllImport("advapi32.dll", ExactSpelling = true, SetLastError = true)]
    private static extern bool OpenProcessToken(IntPtr h, int acc, ref IntPtr phtok);

    [DllImport("advapi32.dll", SetLastError = true)]
    private static extern bool LookupPrivilegeValue(string? host, string name, ref LUID pluid);

    [DllImport("advapi32.dll", ExactSpelling = true, SetLastError = true)]
    private static extern bool AdjustTokenPrivileges(IntPtr htok, bool disall, ref TOKEN_PRIVILEGES newst, int len, IntPtr prev, IntPtr relen);

    [DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
    [return: MarshalAs(UnmanagedType.Bool)]
    private static extern bool CloseHandle(IntPtr hObject);


    [DllImport("user32.dll")]
    private static extern IntPtr GetShellWindow();

    [DllImport("user32.dll", SetLastError = true)]
    private static extern uint GetWindowThreadProcessId(IntPtr hWnd, out uint lpdwProcessId);

    [DllImport("kernel32.dll", SetLastError = true)]
    private static extern IntPtr OpenProcess(ProcessAccessFlags processAccess, bool bInheritHandle, uint processId);

    [DllImport("advapi32.dll", CharSet = CharSet.Auto, SetLastError = true)]
    private static extern bool DuplicateTokenEx(IntPtr hExistingToken, uint dwDesiredAccess, IntPtr lpTokenAttributes, SECURITY_IMPERSONATION_LEVEL impersonationLevel, TOKEN_TYPE tokenType,
        out IntPtr phNewToken);

    [DllImport("advapi32", SetLastError = true, CharSet = CharSet.Unicode)]
    private static extern bool CreateProcessWithTokenW(IntPtr hToken, int dwLogonFlags, string lpApplicationName, string lpCommandLine, int dwCreationFlags, IntPtr lpEnvironment,
        string lpCurrentDirectory, [In] ref STARTUPINFO lpStartupInfo, out PROCESS_INFORMATION lpProcessInformation);

    #endregion
}